package app.chatroom.web.websocket;


import app.chatroom.web.sys_ollama_ai.Service.OllAmaStreamingService;
import jakarta.websocket.OnClose;
import jakarta.websocket.OnMessage;
import jakarta.websocket.OnOpen;
import jakarta.websocket.Session;
import jakarta.websocket.server.PathParam;
import jakarta.websocket.server.ServerEndpoint;


import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.Map;

@Slf4j
@Component
@ServerEndpoint("/ws/{id}")
public class webSocketServer {

    // 注入的时候，给类的 service 注入
    private static OllAmaStreamingService ollAmaStreamingService;


    // 注入的时候，给类的 service 注入
    @Autowired
    public void setOllAmaAiTestService(OllAmaStreamingService ollAmaStreamingService) {
        webSocketServer.ollAmaStreamingService = ollAmaStreamingService;
    }

    // 存储所有的session
    private static final Map<String, Session> sessionMap = new HashMap<>();


    // 链接成功
    @OnOpen
    public void onOpen(Session session, @PathParam("id") String id) throws InterruptedException {
        System.out.println("链接成功");
        System.out.println("客户端id:" + id);
        sessionMap.put(id, session);
        // 发送消息
        sendMessage(ollAmaStreamingService.hello(), id);
    }

    // 接收消息
    @OnMessage
    public void onMessage(String message, @PathParam("id") String id) {

//        String s = ollAmaStreamingService.callTxt(message);
//        sendMessage(s, id);

        log.info("收到消息：{} \r\n id: {}", message, id);

        Flux<String> stringFlux = ollAmaStreamingService.streamGeneration(message);
        StreamMessages(stringFlux, id);

    }

    // 链接关闭
    @OnClose
    public void onClose(@PathParam("id") String id) {
        Session remove = sessionMap.remove(id);
        System.out.println(id + "链接关闭,删除：" + remove);
    }


    // 发送消息
    public void sendMessage(String message, String id) throws InterruptedException {
        System.out.println("发送消息:" + message);
        Session session = sessionMap.get(id);
        String[] split = { message, "[|]"};
        for (String s : split) {
            session.getAsyncRemote().sendText(s);
            Thread.sleep(10);
        }

    }

    //流式发送消息
    // 使用 sessionMap.get(id).getAsyncRemote().sendText("数据") 发数据
    // 使用 [DONE] 结束
    // 帮我完成

    /**
     * 流式发送消息方法
     *
     * @param message 消息流
     * @param id      会话标识符
     */
    public void StreamMessages(Flux<String> message, String id) {
        Session session = sessionMap.get(id);
        if (session == null) {
            log.error("找不到指定ID的会话：{}", id);
            return;
        }

        // 使用doOnNext处理每个元素
        message.doOnNext(msg -> {
            // 发送文本消息
            session.getAsyncRemote().sendText(msg);
        }).doOnComplete(() -> {
            // 最后发送[DONE]标志
            session.getAsyncRemote().sendText("[|]");
        }).subscribe();
    }
    // 群发送消息
    //    public void sendMessage(String message) {
    //        System.out.println("群发送消息:" + message);
    //        for (Session value : sessionMap.values()) {
    //            try {
    //                value.getAsyncRemote().sendText(message);
    //            } catch (Exception e) {
    //                e.printStackTrace();
    //            }
    //        }
    //    }
}
